import random
import torch
import d2l
from chinese_split_word_ictclas_custom import load_corpus_time_machine


def seq_data_iter_random(corpus, batch_size, num_steps):  # @save
    """使用随机抽样生成一个小批量子序列"""
    # print(corpus)
    # 从随机偏移量开始对序列进行分区，随机范围包括num_steps-1
    corpus = corpus[random.randint(0, num_steps - 1):]
    # print(corpus)
    # 减去1，是因为我们需要考虑标签
    num_subseqs = (len(corpus) - 1) // num_steps
    # print("num_subseqs",num_subseqs)

    # 长度为num_steps的子序列的起始索引
    initial_indices = list(range(0, num_subseqs * num_steps, num_steps))
    # 在随机抽样的迭代过程中，
    # 来自两个相邻的、随机的、小批量中的子序列不一定在原始序列上相邻
    random.shuffle(initial_indices)

    # print("initial_indices",initial_indices)

    def data(pos):
        # 返回从pos位置开始的长度为num_steps的序列
        # print("pos",pos)
        return corpus[pos: pos + num_steps]

    # 6//2=3
    num_batches = num_subseqs // batch_size
    for i in range(0, batch_size * num_batches, batch_size):
        # print(i)
        # 在这里，initial_indices包含子序列的随机起始索引 0,2 2,4 4,6
        initial_indices_per_batch = initial_indices[i: i + batch_size]
        X = [data(j) for j in initial_indices_per_batch]
        Y = [data(j + 1) for j in initial_indices_per_batch]
        yield torch.tensor(X), torch.tensor(Y)


# my_seq = list(range(30))
# for X, Y in seq_data_iter_random(my_seq, batch_size=1, num_steps=5):
#     print('X: ', X)
#     # print('X: ', X, '\nY:', Y) n

# 8.3.4.2. 顺序分区

def seq_data_iter_sequential(corpus, batch_size, num_steps):  # @save
    """使用顺序分区生成一个小批量子序列"""
    # 从随机偏移量开始划分序列
    offset = random.randint(0, num_steps)
    num_tokens = ((len(corpus) - offset - 1) // batch_size) * batch_size
    Xs = torch.tensor(corpus[offset: offset + num_tokens])
    Ys = torch.tensor(corpus[offset + 1: offset + 1 + num_tokens])
    Xs, Ys = Xs.reshape(batch_size, -1), Ys.reshape(batch_size, -1)
    num_batches = Xs.shape[1] // num_steps
    for i in range(0, num_steps * num_batches, num_steps):
        X = Xs[:, i: i + num_steps]
        Y = Ys[:, i: i + num_steps]
        yield X, Y


# print("顺序分区")
# for X, Y in seq_data_iter_sequential(my_seq, batch_size=2, num_steps=5):
#     print('X: ', X, '\nY:', Y)


class SeqDataLoader:  # @save
    """加载序列数据的迭代器"""

    def __init__(self, batch_size, num_steps, use_random_iter):
        if use_random_iter:
            self.data_iter_fn = seq_data_iter_random
        else:
            self.data_iter_fn = seq_data_iter_sequential
        self.corpus, self.vocab, _ = load_corpus_time_machine()
        # ##################################
        # print("##########################")
        # #  词指向id
        # print(self.vocab["number"])
        # # id指向词语
        # print(self.vocab.to_tokens(0))
        # print(self.vocab.to_tokens([0, 1, 2]))
        # print("##########################")
        # ##################################

        # self.corpus 是一个长列表，每个元素是一个字符的索引
        # print("111111111",len(self.corpus))
        self.batch_size, self.num_steps = batch_size, num_steps

    def __iter__(self):
        return self.data_iter_fn(self.corpus, self.batch_size, self.num_steps)


def load_data_time_machine(batch_size, num_steps,  # @save
                           use_random_iter=False):
    data_iter = SeqDataLoader(
        batch_size, num_steps, use_random_iter)
    return data_iter, data_iter.vocab

# data_iter, vocab = load_data_time_machine(16, 5, True)
# print(len(vocab))
# # print("顺序分区")
# # 对应图8.4.2 基于循环神经网络的字符级语言模型
# for X, Y in data_iter:
#     print('X: ', X, '\nY:', Y)

# # print("##########################")
# # #  词指向id
# # print(vocab["number"])
# # # id指向词语
# # print(vocab.to_tokens(0))
# # print(vocab.to_tokens([0, 1, 2]))
# # print("##########################")
